#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Time    : 2021/8/1 1:24 上午
# @Author  : WangZhixing

import torch
from torch_geometric.nn import GraphConv
import torch.nn.functional as F

class GraphSage(torch.nn.Module):
    def __init__(self, num_features,hidden_channels,num_classes):
        super(GraphSage, self).__init__()
        self.conv1 = GraphConv(num_features, hidden_channels)
        self.conv2 = GraphConv(hidden_channels, num_classes)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = x.relu()

        # x = x.sigmoid()
        # x = x.tanh()

        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        return x